import argparse


def args_parser():
    """Parse command line arguments for federated learning configuration."""
    parser = argparse.ArgumentParser()

    # Federated learning parameters
    parser.add_argument('--epochs', type=int, default=100, help="Total communication rounds")
    parser.add_argument('--num_users', type=int, default=10, help="Number of clients")
    parser.add_argument('--frac', type=float, default=1.0, help="Fraction of participating clients")
    parser.add_argument('--local_ep', type=int, default=2, help="Local training epochs")
    parser.add_argument('--local_bs', type=int, default=64, help="Local batch size")
    parser.add_argument('--bs', type=int, default=128, help="Test batch size")
    parser.add_argument('--lr', type=float, default=0.05, help="Learning rate")
    parser.add_argument('--momentum', type=float, default=0.5, help="SGD momentum")
    parser.add_argument('--split', type=str, default='user', help="Train-test split type")
    parser.add_argument('--lr_decay', type=float, default=0.975, help="Learning rate decay rate")

    # Model parameters
    parser.add_argument('--model', type=str, default='mlp', help="Model name")
    parser.add_argument('--kernel_num', type=int, default=9, help="Number of each kernel type")
    parser.add_argument('--kernel_sizes', type=str, default='3,4,5', help="Comma-separated kernel sizes")
    parser.add_argument('--norm', type=str, default='batch_norm', help="Normalization type")
    parser.add_argument('--num_filters', type=int, default=32, help="Number of CNN filters")
    parser.add_argument('--max_pool', type=str, default='True', help="Use max pooling")

    # Dataset parameters
    parser.add_argument('--dataset', type=str, default='mnist', help="Dataset name")
    parser.add_argument('--iid', action='store_true', default=True, help="IID data distribution")
    parser.add_argument('--num_classes', type=int, default=10, help="Number of classes")
    parser.add_argument('--num_channels', type=int, default=1, help="Number of image channels")

    # System parameters
    parser.add_argument('--gpu', type=int, default=0, help="GPU ID (-1 for CPU)")
    parser.add_argument('--stopping_rounds', type=int, default=10, help="Early stopping rounds")
    parser.add_argument('--verbose', action='store_true', help="Verbose output")
    parser.add_argument('--seed', type=int, default=None, help="Random seed")
    parser.add_argument('--all_clients', action='store_true', help="Aggregate all clients")

    return parser.parse_args()